// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <numeric>

#include <nebula/compute/exec.h>
#include <nebula/compute/row/grouper.h>
#include <nebula/testing/ktest_util.h>
#include <nebula/testing/random.h>

namespace nebula {
namespace compute {

// Specialized case for GH-40997
TEST(Grouper, ResortedColumnsWithLargeNullRows) {
  const uint64_t num_rows = 1024;

  // construct random array with plenty of null values
  const int32_t kSeed = 42;
  const int32_t min = 0;
  const int32_t max = 100;
  const double null_probability = 0.3;
  const double true_probability = 0.5;
  auto rng = random::RandomArrayGenerator(kSeed);
  auto b_arr = rng.Boolean(num_rows, true_probability, null_probability);
  auto i32_arr = rng.Int32(num_rows, min, max, null_probability);
  auto i64_arr = rng.Int64(num_rows, min, max * 10, null_probability);

  // construct batches with columns which will be resorted in the grouper make
  std::vector<ExecBatch> exec_batches = {ExecBatch({i64_arr, i32_arr, b_arr}, num_rows),
                                         ExecBatch({i32_arr, i64_arr, b_arr}, num_rows),
                                         ExecBatch({i64_arr, b_arr, i32_arr}, num_rows),
                                         ExecBatch({i32_arr, b_arr, i64_arr}, num_rows),
                                         ExecBatch({b_arr, i32_arr, i64_arr}, num_rows),
                                         ExecBatch({b_arr, i64_arr, i32_arr}, num_rows)};

  const int num_batches = static_cast<int>(exec_batches.size());
  std::vector<uint32_t> group_num_vec;
  group_num_vec.reserve(num_batches);

  for (const auto& exec_batch : exec_batches) {
    ExecSpan span(exec_batch);
    ASSERT_OK_AND_ASSIGN(auto grouper, Grouper::create(span.get_types()));
    ASSERT_OK_AND_ASSIGN(Datum group_ids, grouper->consume(span));
    group_num_vec.emplace_back(grouper->num_groups());
  }

  for (int i = 1; i < num_batches; i++) {
    ASSERT_EQ(group_num_vec[i - 1], group_num_vec[i]);
  }
}

}  // namespace compute
}  // namespace nebula
